package com.huangkailong.sqlprinter.agent.mysql;

import com.huangkailong.sqlprinter.agent.mysql.driver.DriverAgent;
import com.huangkailong.sqlprinter.agent.mysql.driver.NewDriverAgent;
import com.huangkailong.sqlprinter.agent.mysql.driver.OldDriverAgent;
import javassist.*;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * @author huangkailong
 * @since 1.0.0
 */
public class MysqlSqlPrinterAgent implements ClassFileTransformer {
    private static final List<DriverAgent> driverAgents;

    static {
        driverAgents = new ArrayList<>();
        driverAgents.add(new NewDriverAgent());
        driverAgents.add(new OldDriverAgent());
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        try {
            if(Objects.nonNull(className) && className.length() > 0){
                final String newClassName = className.replaceAll("/", ".");
                for (DriverAgent driverAgent : driverAgents) {
                    if(driverAgent.isSupport(newClassName)){
                        final CtClass ctClass = driverAgent.enhanceClass();
                        if(Objects.nonNull(ctClass)){
                            System.out.println(newClassName+" 增强成功!");
                            return ctClass.toBytecode();
                        }
                    }
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static List<DriverAgent> getDriverAgents(){
        return driverAgents;
    }
}
